// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <string>
#include <vector>

#include <tbir/runtime/container.h>
#include <tbir/runtime/logging.h>
#include <tbir/runtime/string_util.h>


namespace tbir::runtime::internal {
    inline int cuda_device_offset(int gpu_ordinal) {
        if (gpu_ordinal < 0) {
            return gpu_ordinal;
        }
        // use `index` instead of `device.oridinal` for pytorch engine
        // e.g. if CUDA_VISIBLE_DEVICES=5,6,7 here should be 0,1,2
        char const *tmp = getenv("CUDA_VISIBLE_DEVICES");
        if (tmp == nullptr) {
            // env not set, use all gpus
            return gpu_ordinal;
        }

        std::string env_val(tmp);
        std::vector<std::string> gpu_list = StringUtil::Split(env_val, ",");
        auto it = std::find(gpu_list.begin(), gpu_list.end(), std::to_string(gpu_ordinal));
        MXCHECK(it != gpu_list.end()) << "gpu_ordinal:" << gpu_ordinal
                                      << " not found in CUDA_VISIBLE_DEVICES:" << env_val;
        return std::distance(gpu_list.begin(), it);
    }

    inline int get_tf_inter_op_threads() {
        char const *tmp = getenv("TF_INTER_OP_THREADS");
        if (tmp == nullptr) {
            // env not set, use default 1
            return 1;
        }
        try {
            return std::stoi(std::string(tmp));
        } catch (...) {
            return 1;
        }
    }

    inline int get_tf_intra_op_threads() {
        char const *tmp = getenv("TF_INTRA_OP_THREADS");
        if (tmp == nullptr) {
            // env not set, use default 1
            return 1;
        }
        try {
            return std::stoi(std::string(tmp));
        } catch (...) {
            return 1;
        }
    }

    inline void ListToStdVector_String(const List &list, std::vector<std::string> &result) {
        result.reserve(list.size());
        for (auto &item : list) {
            switch (item.type_code()) {
                case TypeIndex::kRuntimeString: {
                    auto item_unwrapped = item.As<String>();
                    result.push_back(std::move(item_unwrapped));
                }
                    break;
                case TypeIndex::kRuntimeUnicode: {
                    auto item_unwrapped = item.As<Unicode>().encode();
                    result.push_back(std::move(item_unwrapped));
                }
                    break;
                default: {
                    /* not compatible type */
                    MXCHECK(false) << "input type error, \n"
                                   << "optional: str, \n"
                                   << "but receive type : " << item.type_name();
                }
                    break;
            }
        }
    }

}  // namespace tbir::runtime::internal
